import numpy as np
from PyQt6.QtWidgets import QFileDialog, QInputDialog

import pyqtgraph as pg
from pyqtgraph.Qt import QtCore
# from pyqtgraph.GraphicsScene.mouseEvents import MouseClickEvent, MouseDragEvent

class CustomEditedGraphItem(pg.GraphItem):
    """
    可以用于编辑的节点图。
    1，读取数据。
    2，保存数据。
    3，添加节点。
    4，建立连接。
    5，断开连接。
    6，设置点的坐标。
    7，鼠标拖动节点坐标。
    """
    def __init__(self, file_path: str=None, **kwds):
        """
        初始化函数。
        :param file_path: 保存节点坐标和连接关系的json文件，有默认数据, defaults to None
        """
        from PyQt6.QtWidgets import (QMenu)
        self.__dragPoint__ = None  # 当前拖拽的点。
        self.__dragOffset__ = None  # 拖拽偏移量。
        self.textItems = []
        self.text = None
        pg.GraphItem.__init__(self)
        self.context_menu = QMenu()  # 右键菜单。
        if file_path is None:
            # 添加默认数据。
            pos = np.array([[0,0], [10,0]], dtype=float)  # 定义节点。
            adj = np.array([[0,1]])  # 定义图的连接。
            lines = np.array([
                        (255, 0, 0, 255, 1)] * adj.shape[0],
                        dtype=[('red',np.ubyte), ('green',np.ubyte), ('blue',np.ubyte),
                            ('alpha',np.ubyte), ('width',float)])
            texts = ["Point %d" % i for i in range(pos.shape[0])]  # 定义每个点的文本。
            self.setData(pos=pos, adj=adj, size=1, pxMode=False, text=texts, pen=lines)  # 更新图。
    
    def read_data(self):
        """
        读取json数据。包括自动定义文本。
        """
        import json
        file_path, _ = QFileDialog.getOpenFileName(directory='./', filter='JSON Files (*.json)')
        if file_path:
            with open(file_path, 'r') as f:
                data = json.load(f)  # 读取JSON文件内容。
                pos = np.array(data['pos'])  # 将点的坐标转换为numpy数组。
                adj = np.array(data['adj'], dtype=np.int32)  # 将连接关系转换为numpy数组。
                lines = np.array([(255, 0, 0, 255, 1)] * adj.shape[0],
                                dtype=[('red', np.ubyte), ('green', np.ubyte), ('blue', np.ubyte),
                                    ('alpha', np.ubyte), ('width', float)])
                texts = ["Point %d" % i for i in range(pos.shape[0])]  # 定义每个点的文本。
                self.setData(pos=pos, adj=adj, size=1, pxMode=False, text=texts, pen=lines)

    def save_data(self):
        """
        保存当前所有点的坐标和连接关系到json文件中。
        """
        data = {}
        xs, ys = self.scatter.getData()  # 返回点的坐标。
        data['pos'] = np.hstack((xs.reshape(-1, 1), ys.reshape(-1, 1))).tolist()
        data['adj'] = self.adjacency.tolist()
        # print(data)
        file_path, _ = QFileDialog.getSaveFileName(directory='./', filter='JSON Files (*.json)')
        # 添加以下代码来实际保存文件。
        if file_path:  # 确保用户没有取消保存对话框。
            with open(file_path, 'w') as f:
                # 手动构建JSON字符串，确保格式符合要求。
                f.write('{\n')
                # 写入pos键。
                f.write('    "pos": [\n')
                pos_list = data['pos']
                for i, (x, y) in enumerate(pos_list):
                    f.write(f'        [{x: .6f}, {y: .6f}]')  # 保留6位小数。
                    if i != len(pos_list) - 1:
                        f.write(',')
                    f.write('\n')
                f.write('    ],\n')
                # 写入adj键。
                f.write('    "adj": [\n')
                pos_list = data['adj']
                for row, (i, j) in enumerate(pos_list):
                    f.write(f'        [{i}, {j}]')  # 保留6位小数。
                    if row != len(pos_list) - 1:
                        f.write(',')
                    f.write('\n')
                f.write('    ]\n')

                f.write('}')

    def add_point_at_pos(self, pos):
        """
        在指定位置添加一个新点。
        :param pos: QPointF,局部坐标。
        """
        # 获取当前数据
        current_pos = self.data['pos']
        current_text = self.text
        new_pos = np.array([pos.x(), pos.y()])  # 计算新点的位置(将QPointF转换为numpy数组)。
        self.data['pos'] = np.vstack([current_pos, new_pos])  # 添加新点到pos数组。
        # 添加新点的文本。
        new_text = f"Point {len(current_text)}"
        self.text = current_text + [new_text]
        self.setData(pos=self.data['pos'], adj=self.data['adj'], text=self.text,
                    size=1, pxMode=False, pen=self.data['pen'])  # 更新数据。

    def setData(self, **kwds):
        """
        （重写父类方法。）重设点的坐标和建立连接，并更新图。
        """
        self.text = kwds.pop('text', [])
        self.data = kwds
        if 'pos' in self.data:
            npts = self.data['pos'].shape[0]  # 点的数量。
            self.data['data'] = np.empty(npts, dtype=[('index', int)])
            self.data['data']['index'] = np.arange(npts)
        self.setTexts(self.text)
        self.updateGraph()  # 更新图。
        
    def setTexts(self, texts: list[str]):
        """
        设置文本，添加TextItem。
        :param texts: 为所有的点提供的新的文本pg.TextItem。
        """
        # 清除旧的文本项。
        for i in self.textItems:
            i.scene().removeItem(i)
        self.textItems = []
        # 设置新的文本项，并配置颜色和字体。
        for t in texts:
            item = pg.TextItem(t, color=(0, 255, 0, 255), anchor=(0.5, 0.5))  # 绿色，RGBA
            item.setFont(pg.QtGui.QFont("Times New Roman", 10))  # 设置字体和大小，比如 Times New Roman 10px
            self.textItems.append(item)
            item.setParentItem(self)
        
    def updateGraph(self):
        """
        更新图，并设置对应的文本。
        """
        pg.GraphItem.setData(self, **self.data)
        for i,item in enumerate(self.textItems):
            item.setPos(*self.data['pos'][i])

    def connect_point_to_others(self, point_index: int):
        """
        建立连接关系。
        :param point_index: 指定点的序号。
        """
        from PyQt6.QtWidgets import QInputDialog
        available_points = {i for i in range(self.data['pos'].shape[0]) if i != point_index}
        # 检查已经被连接的点。
        adj: np.ndarray = self.data['adj']
        adj_list: list = adj.tolist()
        connected_points = set()
        for edge in adj_list:
            if edge[0] == point_index:
                connected_points.add(edge[1])
            elif edge[1] == point_index:
                connected_points.add(edge[0])
        available_points = available_points.difference(connected_points)  # 过滤掉已经连接的点。
        if not available_points:
            print(f"点 {point_index} 没有可连接的点。")
            return

        point_options = [p for p in available_points]
        point_options.sort()
        point_options = [f"Point {p}" for p in point_options]  # 创建选项文本列表。
        selected_point, ok = QInputDialog.getItem(
            None, "连接到点", f"选择要连接到点 {point_index} 的点:", point_options, 
            current=0, editable=False)  # 弹出选择框，默认选择第一个选项。
        if ok and selected_point:
            selected_index = int(selected_point.split(" ")[1])  # 提取点索引(从"Point X"中提取X)。
            # 建立连接。
            if point_index > selected_index:
                point_index, selected_index = selected_index, point_index  # 调换成由小到大的顺序。
            new_connection = np.array([point_index, selected_index])
            self.data['adj'] = np.vstack([self.data['adj'], new_connection])
            self.data['pen'] = np.array([(255, 0, 0, 255, 1)] * self.data['adj'].shape[0],
                                        dtype=[('red', np.ubyte), ('green', np.ubyte), ('blue', np.ubyte),
                                                ('alpha', np.ubyte), ('width', float)])
            self.updateGraph()
        
    def disconnect_point(self, point_index: int):
        """
        断开指定点的所有连接。
        :param point_index: 要断开连接的点的索引。
        """
        adj: np.ndarray = self.data['adj']
        adj_list: list = adj.tolist()
        # 找到所有包含 point_index 的边。
        connected_edges = [edge for edge in adj_list if point_index in edge]
        if not connected_edges:
            print(f"点 {point_index} 没有可断开的连接。")
            return
        # 创建选项文本列表，显示与 point_index 相连的点。
        edge_options = []
        for edge in connected_edges:
            other_point = edge[1] if edge[0] == point_index else edge[0]
            edge_options.append(other_point)
        edge_options.sort()
        edge_options = [f"Point {item}" for item in edge_options]
        selected_edge, ok = QInputDialog.getItem(
            None, "断开连接", f"选择要断开点 {point_index} 的连接:", edge_options, 
            current=0, editable=False)  # 弹出选择框，默认选择第一个选项。
        
        if ok and selected_edge:
            # 提取要断开的边的索引。
            other_point = int(selected_edge.split(" ")[1])  # 提取其他点的索引。
            edge_to_remove = None
            for edge in adj_list:
                if edge == [point_index, other_point] or edge == [other_point, point_index]:
                    edge_to_remove = edge
                    break
            if edge_to_remove is not None:
                # 从邻接矩阵中移除该边。
                new_adj = [e for e in adj_list if e != edge_to_remove]
                self.data['adj'] = np.array(new_adj, dtype=np.int32).reshape([-1, 2])  # 防止graph无连接的情况。
                self.data['pen'] = np.array([(255, 0, 0, 255, 1)] * len(new_adj),
                                            dtype=[('red', np.ubyte), ('green', np.ubyte), ('blue', np.ubyte),
                                                    ('alpha', np.ubyte), ('width', float)])
                self.updateGraph()
    
    def set_coordinate(self, point_index: int):
        """
        设置点的坐标。
        :param point_index: 要设置的点的索引。
        """
        text = f"{self.data['pos'][point_index, 0]} {self.data['pos'][point_index, 1]}"
        text, ok = QInputDialog.getText(None, "输入二维坐标", 
                                            "请输入两个数，用空格隔开（默认输入当前节点坐标）:", text=text)
        if not ok:
            return  # 用户点击取消。
        try:
            x_str, y_str = text.split(' ', 1)  # 仅分割一次，避免多个逗号的情况。
            coordinates = float(x_str.strip()), float(y_str.strip())
        except ValueError:
            return  # 非数字输入转换失败，输入的数字不够。
        self.data['pos'][point_index] = np.array(coordinates)  # 获取当前数据。
        self.setData(pos=self.data['pos'], adj=self.data['adj'], text=self.text,
                    size=1, pxMode=False, pen=self.data['pen'])  # 更新数据。

    def mouseClickEvent(self, event):
        """
        鼠标点击事件。鼠标右键点击点可以触发菜单。
        1，读取数据功能。
        2，保存数据功能。
        3，在鼠标坐标处添加点。
        4，右击某个点后，可以建立该点与某个点的连接。
        5，右击某个点后，可以断开该点与某个点的连接。
        6，右击某个点后，可以设置点的坐标。
        :param event: pg.GraphicsScene.mouseEvents.MouseClickEvent
        """
        from PyQt6.QtGui import QAction
        if event.button() == QtCore.Qt.MouseButton.RightButton:
            pos = event.pos()  # 局部坐标（相对于 GraphItem）。
            pts = self.scatter.pointsAt(pos)
            scene_pos = self.mapToScene(pos)  # 转换为场景坐标。
            view: pg.GraphicsView = self.getViewBox().scene().views()[0]  # 获取 ViewBox 关联的 QGraphicsView。
            view_pos = view.mapFromScene(scene_pos)  # 场景坐标 → 视图坐标
            global_pos = view.mapToGlobal(view_pos)  # 视图坐标 → 全局屏幕坐标

            self.context_menu.clear()  # 添加菜单及其内容。
            # 1，读取数据功能。
            a = QAction(text='读取数据', parent=self.context_menu)
            a.triggered.connect(self.read_data)
            self.context_menu.addAction(a)
            # 2，保存数据功能。
            a = QAction(text='保存数据', parent=self.context_menu)
            a.triggered.connect(self.save_data)
            self.context_menu.addAction(a)
            # 3，在鼠标坐标处添加点。
            a = QAction(text='添加节点', parent=self.context_menu)
            a.triggered.connect(lambda: self.add_point_at_pos(pos))
            self.context_menu.addAction(a)
            if len(pts) > 0:
                # 如果点击到了点。
                clicked_point = pts[0]
                point_index = clicked_point.data()[0]  # 获取被点击点的索引。
                # 4，右击某个点后，可以建立该点与某个点的连接。
                a = QAction(text='建立连接', parent=self.context_menu)
                a.triggered.connect(lambda: self.connect_point_to_others(point_index))
                self.context_menu.addAction(a)
                # 5，右击某个点后，可以断开与其他所有点的连接。
                a = QAction(text='断开连接', parent=self.context_menu)
                a.triggered.connect(lambda: self.disconnect_point(point_index))
                self.context_menu.addAction(a)
                # 6，右击某个点后，可以设置点的坐标。
                a = QAction(text='设置坐标', parent=self.context_menu)
                a.triggered.connect(lambda: self.set_coordinate(point_index))
                self.context_menu.addAction(a)
            self.context_menu.exec(global_pos)  # 在鼠标位置生成菜单并等待。
        event.accept()

    def mouseDragEvent(self, event):
        """
        鼠标左键拖动事件。及时更新点的坐标。
        :param event: pg.GraphicsScene.mouseEvents.MouseDragEvent
        """
        if event.button() != QtCore.Qt.MouseButton.LeftButton:
            event.ignore()
            return
        if event.isStart():
            # 按钮第一次被按下时，找到该点。
            pos = event.buttonDownPos()  # 鼠标的坐标。
            pts = self.scatter.pointsAt(pos)  # 需要获取点的序号。
            if len(pts) == 0:
                event.ignore()  # 没有按下点。
                return
            self.__dragPoint__ = pts[0]
            ind = pts[0].data()[0]
            self.__dragOffset__ = self.data['pos'][ind] - pos
        elif event.isFinish():
            self.__dragPoint__ = None
            return
        else:
            if self.__dragPoint__ is None:
                event.ignore()
                return
        # 更新拖动点的位置信息。
        ind = self.__dragPoint__.data()[0]
        self.data['pos'][ind] = event.pos() + self.__dragOffset__
        self.updateGraph()
        event.accept()

pg.setConfigOptions(antialias=True)

if __name__ == '__main__':
    w = pg.GraphicsLayoutWidget(show=True)
    w.setWindowTitle('Graph关系')
    v: pg.ViewBox = w.addViewBox()
    v.setAspectLocked()
    g = CustomEditedGraphItem()
    v.addItem(g)
    v.setMenuEnabled(False)
    pg.exec()
